import psycopg2
import re

class OpenGaussSQLExecutor:
    def __init__(self, host, port, database, user, password):
        """
        初始化数据库连接参数
        :param host: 数据库主机地址
        :param port: 数据库端口
        :param database: 数据库名称
        :param user: 用户名
        :param password: 密码
        """
        self.conn_params = {
            'host': host,
            'port': port,
            'database': database,
            'user': user,
            'password': password
        }
        self.connection = None
        self.cursor = None

    def connect(self):
        """建立到 openGauss 的数据库连接"""
        try:
            self.connection = psycopg2.connect(**self.conn_params)
            self.cursor = self.connection.cursor()
            print("Successfully connected to openGauss database")
        except Exception as e:
            print(f"Connection failed: {str(e)}")
            raise

    def disconnect(self):
        """关闭数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
        print("Database connection closed")

    def _parse_sql_scripts(self, file_path):
        """
        解析 SQL 文件，分割成独立的 SQL 语句
        :param file_path: SQL 文件路径
        :return: 包含独立 SQL 语句的列表
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            sql_content = f.read()
        
        # 移除注释（支持 -- 和 /* */ 两种格式）
        sql_content = re.sub(r'/\*.*?\*/', '', sql_content, flags=re.DOTALL)
        sql_content = re.sub(r'--.*', '', sql_content)
        
        # 按分号分割并清理空白语句
        statements = [stmt.strip() for stmt in sql_content.split(';') if stmt.strip()]
        return statements

    def execute_sql_file(self, file_path, autocommit=False):
        """
        执行 SQL 文件中的所有脚本
        :param file_path: SQL 文件路径
        :param autocommit: 是否自动提交事务（默认为 False）
        """
        try:
            self.connect()
            statements = self._parse_sql_scripts(file_path)
            
            for i, statement in enumerate(statements, 1):
                try:
                    print(f"Executing statement {i}/{len(statements)}")
                    self.cursor.execute(statement)
                    if autocommit:
                        self.connection.commit()
                except Exception as e:
                    print(f"Error executing statement {i}: {str(e)}")
                    print(f"Failed statement: {statement[:100]}...")  # 打印前100字符
                    if not autocommit:
                        self.connection.rollback()
                    raise
            
            if not autocommit:
                self.connection.commit()
            print("All SQL statements executed successfully")
            
        except Exception as e:
            print(f"Execution failed: {str(e)}")
            if self.connection and not autocommit:
                self.connection.rollback()
            raise
        finally:
            self.disconnect()

# 使用示例
if __name__ == "__main__":
    # 配置数据库连接参数
    executor = OpenGaussSQLExecutor(
        host="127.0.0.1",
        port="5432",
        database="mydb",
        user="myuser",
        password="mypassword"
    )
    
    # 执行 SQL 文件
    executor.execute_sql_file(
        file_path="scripts.sql", 
        autocommit=False  # 设置为 True 时每个语句独立提交
    )
